'''
@author Michael J Bommarito II
@date Jul 14, 2010
'''

import cairo
import igraph
import math

from graphSequence import graphSequence

class render2D(object):
    '''
    Renderer that produces a 2D animation of a graphSequence object.
    '''

    def __init__(self, sequence = None):
        '''
        Default constructor.
        '''
        
        
        '''
        Animation settings.
        '''
        # Frame counter
        self.frame = 0
        
        # Frame sequence
        self.frameSequence = [] 

        # Frames per second in resulting animation
        self.framesPerSecond = 10
        
        # Number of interpolating frames between layouts
        self.framesInterpolate = 30
        
        # Number of delay frames at the end of each layout
        self.framesDelay = 15
        
        # Number of delay frames at the end of the animation
        self.framesFinal = 30
        
        # Width in pixels
        self.width = 1600
        
        # Height in pixels
        self.height = 1200
        
        # Vertex position vertexLabel -> t -> (x,y)
        self.vertexPosition = {}
        
        # Scaled vertex position
        self.vertexScaledPosition = {}
        
        '''
        Formatting settings
        '''
        # Background color
        self.colorBG = (0.0, 0.0, 0.0)
        
        # Default vertex color
        self.colorVertex = (1.0, 0.0, 0.0)
        
        # Default edge color
        self.colorEdge = (1.0, 1.0, 1.0)
        
        # Default vertex text color
        self.colorTextVertex = (1.0, 1.0, 1.0)
        
        # Default edge text color
        self.colorTextEdge = (1.0, 1.0, 1.0)
        
        # Do we want curved edges?
        self.curvedEdges = True
        
        '''
        If a sequence was passed, check it.
        '''
        if sequence:
            self.graphSequence = sequence
        else:
            self.graphSequence = None
        
        if len(self.graphSequence) > 0:
            self.checkGraphSequence()

    def addGraph(self, graph, label = None):
        '''
        Add a single graph element to the sequence.
        '''
        if not isinstance(graph, igraph.Graph):
            raise Exception("graphmovie::render2D::addGraph: graph is not an igraph.Graph object.")
        
        if graph.vcount() == 0:
            raise Exception("graphmovie::render2D::addGraph: graph has 0 vertices.")
        
        self.graphSequence.addGraph(graph, label)
    
    def checkGraphSequence(self):
        '''
        Check to make sure that our current graph sequence is OK.
        Every element must be an igraph object with at least 1 vertex.
        '''
        
        for i, graph in enumerate(self.graphSequence.graphs):
            if not isinstance(graph, igraph.Graph):
                raise Exception("graphmovie::render2D::checkGraphSequence: element {0} is not an igraph.Graph object.".format(i))
            
            if graph.vcount() == 0:
                raise Exception("graphmovie::render2D::checkGraphSequence: element {0} has 0 vertices.".format(i))

    def getVertexPosition(self, label, t):
        '''
        Get a vertex position given its label and time.  
        Note that if the vertex didn't exist at the time, we
        return an entrance coordinate. 
        '''
        if self.vertexPosition.has_key(label) and self.vertexPosition[label].has_key(t):
            return self.vertexPosition[label][t]
        else:
            return (0.0, 0.0)
        
    def getVertexScaledPosition(self, label, t):
        '''
        Get a vertex position given its label and time.  
        Note that if the vertex didn't exist at the time, we
        return an entrance coordinate. 
        '''
        if self.vertexScaledPosition.has_key(label) and self.vertexScaledPosition[label].has_key(t):
            return self.vertexScaledPosition[label][t]
        elif t >= len(self.graphSequence):
            lastTime = max(self.vertexScaledPosition[label].keys())
            return self.vertexScaledPosition[label][lastTime]
        else:
            return (0.5, 0.5)

    def setVertexPosition(self, label, t, x, y):
        '''
        Set a vertex position given its label, the current time, and its coordinates.
        '''
        if self.vertexPosition.has_key(label):
            self.vertexPosition[label][t] = (x,y)
        else:
            self.vertexPosition[label] = {t: (x,y)}
        
    def setVertexScaledPosition(self, label, t, x, y):
        '''
        Set a vertex position given its label, the current time, and its coordinates.
        '''
        if x < 0.1:
            x += 0.15
        elif x >= 0.9:
            x -= 0.15
            
        if y < 0.1:
            y += 0.15
        elif y >= 0.9:
            y -= 0.15

        if self.vertexScaledPosition.has_key(label):
            self.vertexScaledPosition[label][t] = (x,y)
        else:
            self.vertexScaledPosition[label] = {t: (x,y)}

    def getSeedLayout(self, t):
        '''
        Get the seed layout for the graph at time t.
        This requires us to determine the vertices at t
        and then get their position at t-1.
        '''
        graph = self.graphSequence.graphs[t]
        layout = []
        
        for v in graph.vs:
            layout.append(self.getVertexPosition(v['label'], t-1))
        
        return layout

    def calculateLayout(self):
        '''
        Calculate the layout of each vertex for each graph in the sequence.
        In this case, these coordinates are simple (t,x,y).
        '''
        self.calculateInitialLayout()
        
        for t in range(1, len(self.graphSequence)):
            graph = self.graphSequence.graphs[t]
            layout = self.getSeedLayout(t)
            layout = graph.layout_kamada_kawai(seed = layout, maxiter = 2000)
            
            x, y = zip(*layout)
            minX, maxX, minY, maxY = min(x), max(x), min(y), max(y)
            xRange, yRange = maxX - minX, maxY - minY
            
            for i,v in enumerate(graph.vs):
                self.setVertexPosition(v['label'], t, layout[i][0], layout[i][1])
                self.setVertexScaledPosition(v['label'], t, (layout[i][0] - minX) / xRange, (layout[i][1] - minY) / yRange)

    def calculateInitialLayout(self):
        '''
        Calculate the initial layout of each vertex for the first graph.
        This is the seed layout for the remainder of the animation, so we
        need to make sure it is good.
        '''
        graph = self.graphSequence.graphs[0]
        layout = graph.layout_circle()
        layout = graph.layout_kamada_kawai(seed = layout, maxiter = 5000)
        
        x, y = zip(*layout)
        minX, maxX, minY, maxY = min(x), max(x), min(y), max(y)
        xRange, yRange = maxX - minX, maxY - minY
                
        for i, v in enumerate(graph.vs):
            self.setVertexPosition(v['label'], 0, layout[i][0], layout[i][1])
            self.setVertexScaledPosition(v['label'], 0, (layout[i][0] - minX) / xRange, (layout[i][1] - minY) / yRange) 

    def initSurface(self, surf):
        '''
        Initialize a Cairo surface.
        '''
        con = cairo.Context(surf)
        con.scale(self.width / 1.0, self.height / 1.0)
        
        con.set_source_rgba(self.colorBG[0], self.colorBG[1], self.colorBG[2], 1.0)
        con.rectangle(0, 0, 1.0, 1.0)
        con.fill()
                
        return con
    
    def drawVertex(self, con, t, v, n):
        '''
        Draw a vertex onto a Cairo context.
        n is the interpolant frame.
        '''
        x0, y0 = self.getVertexScaledPosition(v['label'], t)
        x1, y1 = self.getVertexScaledPosition(v['label'], t+1)
        
        x = x0 + (x1 - x0) * (float(n) / self.framesInterpolate)
        y = y0 + (y1 - y0) * (float(n) / self.framesInterpolate)
        
        con.move_to(x, y)
        con.arc(x, y, 0.01, 0.0, 2.0 * math.pi)
        
        try:
            con.set_source_rgba(v['graphmovie_color'][0], v['graphmovie_color'][1], v['graphmovie_color'][2], 1.0)
        except:
            con.set_source_rgba(self.colorVertex[0], self.colorVertex[1], self.colorVertex[2], 1.0)

        con.fill()
        
        '''
        Draw the label if this vertex has the 
        graphmovie_label property set.
        ''' 

        if 'graphmovie_label' in v.attribute_names() and v['graphmovie_label']:
            label = v['graphmovie_label']
            if 'graphmovie_label_color' in v.attribute_names() and v['graphmovie_label_color']:
                color = v['graphmovie_label_color']
            else:
                color = self.colorTextVertex
            
            if y < 0.5:
                con.move_to(x, y - 0.011)
            else:
                con.move_to(x, y + 0.011)
            
            con.set_font_size(0.015)
            con.set_source_rgba(color[0], color[1], color[2], 1)
            con.show_text(v['graphmovie_label'])
                
    def drawEdge(self, con, t, e, vA, vB, n):
        '''
        Draw an edge between vA and vB at interpolant n.
        '''
        x0, y0 = self.getVertexScaledPosition(vA['label'], t)
        x1, y1 = self.getVertexScaledPosition(vA['label'], t+1)
        xA = x0 + (x1 - x0) * (float(n) / self.framesInterpolate)
        yA = y0 + (y1 - y0) * (float(n) / self.framesInterpolate)
        
        x0, y0 = self.getVertexScaledPosition(vB['label'], t)
        x1, y1 = self.getVertexScaledPosition(vB['label'], t+1)
        xB = x0 + (x1 - x0) * (float(n) / self.framesInterpolate)
        yB = y0 + (y1 - y0) * (float(n) / self.framesInterpolate)
        
        con.move_to(xA, yA)
        con.set_line_cap(cairo.LINE_JOIN_ROUND)
        con.set_line_width(0.001)
        
        try:
            con.set_source_rgba(e['color'][0], e['color'][1], e['color'][2], 1.0)
        except:
            con.set_source_rgba(self.colorEdge[0], self.colorEdge[1], self.colorEdge[2], 1.0)
        
        if self.curvedEdges:
            c1 = 0.25
            c2 = 0.8
            con.curve_to(xA + (xB - xA) * (yB - yA) * c1, yA + (xB - xA) * (yB - yA) * c1,
                         xA + (xB - xA) * (yB - yA) * c2, yA + (xB - xA) * (yB - yA) * c2,
                         xB, yB)
        else:
            con.line_to(xB, yB)
        
        con.stroke()
        
        '''
        Draw the label if this edge has the 
        graphmovie_label property set.
        ''' 

        if 'graphmovie_label' in e.attribute_names() and e['graphmovie_label']:
            label = e['graphmovie_label']
            if 'graphmovie_label_color' in e.attribute_names() and e['graphmovie_label_color']:
                color = e['graphmovie_label_color']
            else:
                color = self.colorTextEdge
            
            mX = 0.5 * (xA + xB)
            mY = 0.5 * (yA + yB)
            if mY < 0.5:
                con.move_to(mX, mY - 0.011)
            else:
                con.move_to(mX, mY + 0.011)
            
            con.set_font_size(0.015)
            con.set_source_rgba(color[0], color[1], color[2], 1)
            con.show_text(e['graphmovie_label'])
        
    def drawInterpolatedFrames(self, t, graph):
        '''
        Draw the frames for a graph that is not terminal.
        '''
        for n in range(self.framesDelay):
            surf = cairo.ImageSurface(cairo.FORMAT_ARGB32, self.width, self.height)
            con = self.initSurface(surf)
            
            edgelist = graph.get_edgelist()    
            for i, e in enumerate(graph.es):
                self.drawEdge(con, t, e, graph.vs[edgelist[i][0]], graph.vs[edgelist[i][1]], 0)
            
            for v in graph.vs:
                self.drawVertex(con, t, v, 0)

            self.frameSequence.append(surf)

        for n in range(self.framesInterpolate):
            surf = cairo.ImageSurface(cairo.FORMAT_ARGB32, self.width, self.height)
            con = self.initSurface(surf)
            
            edgelist = graph.get_edgelist()    
            for i, e in enumerate(graph.es):
                self.drawEdge(con, t, e, graph.vs[edgelist[i][0]], graph.vs[edgelist[i][1]], n)
            
            for v in graph.vs:
                self.drawVertex(con, t, v, n)

            self.frameSequence.append(surf)

    def drawLastFrame(self, t, graph):
        '''
        Draw the frames for the final graph.
        '''
        for n in range(self.framesFinal):
            surf = cairo.ImageSurface(cairo.FORMAT_ARGB32, self.width, self.height)
            con = self.initSurface(surf)
            
            edgelist = graph.get_edgelist()    
            for i, e in enumerate(graph.es):
                self.drawEdge(con, t, e, graph.vs[edgelist[i][0]], graph.vs[edgelist[i][1]], n)
            
            for v in graph.vs:
                self.drawVertex(con, t, v, n)
            
            self.frameSequence.append(surf)

    def render(self):
        '''
        Render the movie by creating a Cairo surface for each frame.
        '''
        for t, graph in enumerate(self.graphSequence.graphs):
            if t < len(self.graphSequence) - 1:
                self.drawInterpolatedFrames(t, graph)
            else:
                self.drawLastFrame(t, graph)
            
            for surf in self.frameSequence:
                surf.write_to_png('frames/%08d.png' % (self.frame))
                self.frame += 1
            
            self.frameSequence = []
        
            
                
        
        
        
        
        
